using System;
using System.Collections.Generic;
using System.IO;

namespace DeflateLib
{
    /// <summary>
    ///  DEFLATE encoder algorithm implementation.
    /// </summary>
    public sealed class Encoder
    {
        public const int MaxEncodeLength = 0x8000;

        Stream output;

        /// <summary>
        /// Gets the specified output binary stream.
        /// </summary>
        public Stream Output
        {
			get { return output; }
        }
        
        BitWriter writer;

		ICompressor compressor;
        CompressorFactory compressorFactory;

        /// <summary>
        /// Creates encoder that will output data into the specified
        /// binary stream.
        /// </summary>
        /// <param name="output">Output binary stream.</param>
        public Encoder(Stream output)
            : this(output, CompressorFactory.Default)
        {
        }


        /// <summary>
        /// Creates encoder that will output data into the specified
        /// binary stream. The provided compressor will be used to 
        /// compress data.
        /// </summary>
        /// <param name="output">Output binary stream.</param>
        /// <param name="compressorFactory">Copressor factory.</param>
        public Encoder(Stream output, CompressorFactory compressorFactory)
        {
            if (output == null) throw new ArgumentNullException("output");
            if (compressorFactory == null) throw new ArgumentNullException("compressorFactory");

			this.output = output;
			this.writer = new BitWriter(output);
            this.compressorFactory = compressorFactory;

            this.compressor = compressorFactory.CreateCompressor();
            if (this.compressor == null) throw new InvalidOperationException("compressor was not created");
        }

        /// <summary>
        /// Sets dictionary.
        /// </summary>
        /// <param name="data">data</param>
        /// <param name="index">offset with data</param>
        /// <param name="length">data length</param>
        public void SetDictionary(byte[] data, int index, int length)
        {
            ValidateDataIndexLength(data, index, length);

            compressor.Initialize(data, index, length);
        }

        private static void ValidateDataIndexLength(byte[] data, int index, int length)
        {
            if (data == null) throw new ArgumentNullException("data");
            if (index < 0 || index > data.Length) throw new ArgumentOutOfRangeException("index");
            if (length < 0 || length > data.Length) throw new ArgumentOutOfRangeException("length");
            if (index + length > data.Length) throw new ArgumentException("data + length");
        }

		internal void CountCodes(CompressedData compressedData, out int[] codeCounts,
			out int[] distanceCodeCounts, out int extraBitsLength)
		{
            codeCounts = new int[DeflateUtils.MaxCode + 1];
            distanceCodeCounts = new int[DeflateUtils.MaxDistanceCode + 1];
            extraBitsLength = 0;
            int i = 0, j = 0;
            while(i < compressedData.Codes.Count)
            {
				if(compressedData.Codes[i] <= DeflateUtils.EndOfBlockCode) 
				{
					codeCounts[compressedData.Codes[i]]++;
					i++;
				}
				else
				{
					codeCounts[compressedData.Codes[i]]++;
					distanceCodeCounts[compressedData.Codes[i + 2]]++;
					i += 4;
					extraBitsLength += compressedData.CodeBits[j] +
						compressedData.CodeBits[j + 1];
					j += 2;
				}				
            }
		}

        /// <summary>
        /// Write sync block byte-aligned or non byte-aligned.
        /// </summary>
        /// <param name="align">Byte-aligned.</param>
        /// <param name="last">Last block in the stream.</param>
        public void WriteSyncBlock(bool align, bool last)
        {
            writer.WriteBit(last);
            if (align)
            {
                writer.WriteLSB(DeflateUtils.NoCompresstionBlockType, 2);
                writer.Align();
                writer.WriteLSB(0, 16);
                writer.WriteLSB(0xFFFF, 16);
            }
            else
            {
                writer.WriteLSB(DeflateUtils.StandardCodeBlockType, 2);
                TreeNode code = DeflateUtils.StaticCodes[DeflateUtils.EndOfBlockCode];
                writer.WriteMSB(code.Code, code.Len);
            }

            if (last)
            {
                writer.Align();
            }
        }

        /// <summary>
        /// Copy data as is.
        /// </summary>
        /// <param name="data">data</param>
        /// <param name="index">offset with data</param>
        /// <param name="length">data length</param>
        /// <param name="last">Last block in the stream.</param>
        public void Copy(byte[] data, int index, int length, bool last)
        {
            const int MaxDataToCopy = 0xFFFF;
            if (length > MaxDataToCopy) throw new ArgumentOutOfRangeException("length");
            ValidateDataIndexLength(data, index, length);

            writer.WriteBit(last);
            writer.WriteLSB(DeflateUtils.NoCompresstionBlockType, 2);
            writer.Align();
            writer.WriteLSB((UInt16)length, 16);
            writer.WriteLSB((UInt16)~length, 16);
            OutputBytes(data, index, length);
        }

        /// <summary>
        /// Encodes data.
        /// </summary>
        /// <param name="data">data</param>
        /// <param name="index">offset with data</param>
        /// <param name="length">data length</param>
        /// <param name="last">Last block in the stream.</param>
        public void Encode(byte[] data, int index, int length, bool last)
        {
            if (length > MaxEncodeLength) throw new ArgumentOutOfRangeException("length");
            ValidateDataIndexLength(data, index, length);

			CompressedData compressedData = compressor.Compress(data, index, length);
            if (compressedData == null) throw new InvalidOperationException("compressor returned null");
            if (compressedData.LengthWritten != length) throw new InvalidOperationException("compressor did write all data");

            compressedData.WriteEOB();
		
			int[] codeCounts, distanceCodeCounts;
			int extraBitsLength;
			CountCodes(compressedData, out codeCounts, out distanceCodeCounts, out extraBitsLength);

            int[] symbolLengths = DeflateUtils.GetCodeLengths(codeCounts);
            int[] symbolDistanceLengths = DeflateUtils.GetCodeLengths(distanceCodeCounts);

            int totalByPlainCopy = (4 + length) * 8;
            if (writer.BitsToAligment < 3)
                totalByPlainCopy += 5 + writer.BitsToAligment;
            else
                totalByPlainCopy += writer.BitsToAligment - 3;

            int totalWithoutHuffman = DeflateUtils.CalculateLength(symbolLengths, codeCounts) +
                DeflateUtils.CalculateLength(symbolDistanceLengths, distanceCodeCounts) + extraBitsLength;

            int totalInStaticCodes = DeflateUtils.CalculateLengthInStaticCodes(symbolLengths, symbolDistanceLengths)
                + extraBitsLength;

            bool skipDynamicEncoding = totalInStaticCodes <= totalWithoutHuffman
                || totalByPlainCopy <= totalWithoutHuffman;
            if (!skipDynamicEncoding)
            {
                int totalWithHuffman = totalWithoutHuffman + 
                    EstimateHuffmanTree(symbolLengths, symbolDistanceLengths);
                skipDynamicEncoding = totalInStaticCodes <= totalWithHuffman
                    || totalByPlainCopy <= totalWithHuffman;
            }

            bool skipStaticEncoding = totalByPlainCopy < totalInStaticCodes;

            // start writing
            writer.WriteBit(last);
            if (skipStaticEncoding && skipDynamicEncoding)
            {
                writer.WriteLSB(DeflateUtils.NoCompresstionBlockType, 2);
                writer.Align();
                int len = length;
                int nlen = (~length) & 0xFFFF;                
                writer.WriteLSB(len, 16);
                writer.WriteLSB(nlen, 16);
                writer.Align();
                OutputBytes(data, index, length);
            }
            else
            {
                TreeNode[] codes, distanceCodes;

                if (skipDynamicEncoding)
                {
                    writer.WriteLSB(DeflateUtils.StandardCodeBlockType, 2);
                    codes = DeflateUtils.StaticCodes;
                    distanceCodes = DeflateUtils.StaticDistanceCodes;                    
                }
                else
                {
                    writer.WriteLSB(DeflateUtils.DynamicCodeBlockType, 2);

                    WriteHuffmanTree(symbolLengths, symbolDistanceLengths);

                    codes = DeflateUtils.BuildCodes(symbolLengths);
                    distanceCodes = DeflateUtils.BuildCodes(symbolDistanceLengths); 
                }

				WriteCodes(compressedData, codes, distanceCodes);
            }
            
            if (last)
            {
                writer.Align();
            }
        }

        /// <summary>
        /// Clears state such as dictionary, compressor, etc.
        /// </summary>
        public void ClearState()
        {
            this.compressor = compressorFactory.CreateCompressor();
            if (this.compressor == null) throw new InvalidOperationException("compressor was not created");
        }

        private void WriteCodes(CompressedData compressedData, TreeNode[] codesTree, TreeNode[] distanceCodesTree)
        {
            int i = 0, j = 0;
            while(i < compressedData.Codes.Count)
            {
				if(compressedData.Codes[i] <= DeflateUtils.EndOfBlockCode) 
				{
					TreeNode code = codesTree[compressedData.Codes[i]];
					writer.WriteMSB(code.Code, code.Len);
					i++;
				}
				else
				{
					TreeNode code = codesTree[compressedData.Codes[i]];
					writer.WriteMSB(code.Code, code.Len);
					if (compressedData.CodeBits[j] > 0)
					{
						writer.WriteLSB(compressedData.Codes[i + 1], compressedData.CodeBits[j]);
					}
					i += 2;
					j++;

					code = distanceCodesTree[compressedData.Codes[i]];
					writer.WriteMSB(code.Code, code.Len);
					if (compressedData.CodeBits[j] > 0)
					{
						writer.WriteLSB(compressedData.Codes[i + 1], compressedData.CodeBits[j]);
					}
					i += 2;
					j ++;
				}				
            }
        }

        private void OutputBytes(byte[] data, int offset, int length)
        {
            Output.Write(data, offset, length);
        }

        private void WriteHuffmanTree(int[] symbolLengths, int[] symbolDistanceLengths)
        {        
            int[] counts = new int[DeflateUtils.MaxSpecialLengthCode + 1];
            int nonZeroSymbolLengths = DeflateUtils.MaxCode + 1;
            while (nonZeroSymbolLengths > 257 && symbolLengths[nonZeroSymbolLengths - 1] == 0)
                --nonZeroSymbolLengths;

            int nonZeroDistanceLenghts = DeflateUtils.MaxDistanceCode + 1;
            while (nonZeroDistanceLenghts > 1 && symbolDistanceLengths[nonZeroDistanceLenghts - 1] == 0)
                --nonZeroDistanceLenghts;

            int[] lengthsSequence = new int[nonZeroSymbolLengths + nonZeroDistanceLenghts];
            Array.Copy(symbolLengths, 0, lengthsSequence, 0, nonZeroSymbolLengths);
            Array.Copy(symbolDistanceLengths, 0, lengthsSequence, nonZeroSymbolLengths, nonZeroDistanceLenghts);

            int[] lengthsCompact = DeflateUtils.CompactCodeLengths(lengthsSequence, 0, lengthsSequence.Length);
            for (int i = 0; i < lengthsCompact.Length; i++)
            {
                counts[lengthsCompact[i]]++;
                if (lengthsCompact[i] >= DeflateUtils.SpecialLengthCodeStart)
                {
                    ++i;
                }
            }

            const int MaxlengthsCodeLength = 7;
            int[] lengthsCodeLengths = DeflateUtils.GetCodeLengths(counts, MaxlengthsCodeLength);

            int[] map = DeflateUtils.LengthsMap;
            int nonZeroLengthCodes = DeflateUtils.MaxSpecialLengthCode + 1;
            while (nonZeroLengthCodes > 4 && lengthsCodeLengths[map[nonZeroLengthCodes - 1]] == 0)
                --nonZeroLengthCodes;

            int hlit = nonZeroSymbolLengths - 257;
            int hdist = nonZeroDistanceLenghts - 1;
            int hclen = nonZeroLengthCodes - 4;
            TreeNode[] lengthsCodeLengthsTree = DeflateUtils.BuildCodes(lengthsCodeLengths);
            writer.WriteLSB(hlit, 5);
            writer.WriteLSB(hdist, 5);
            writer.WriteLSB(hclen, 4);
            for (int i = 0; i < nonZeroLengthCodes; i++)
            {
                writer.WriteLSB(lengthsCodeLengths[map[i]], 3);
            }

            for (int i = 0; i < lengthsCompact.Length; i++)
            {
                int lengthCode = lengthsCompact[i];
                TreeNode code = lengthsCodeLengthsTree[lengthCode];
                writer.WriteMSB(code.Code, code.Len);
                if (lengthCode >= DeflateUtils.SpecialLengthCodeStart)
                {
                    ++i;
                    writer.WriteLSB(lengthsCompact[i], DeflateUtils.GetExtraBitsForLengthCode(lengthCode));                   
                }
            }

            DeflateUtils.DumpTrees("encoder.log", lengthsCodeLengths, symbolLengths, symbolDistanceLengths);
        }

        private int EstimateHuffmanTree(int[] symbolLengths, int[] symbolDistanceLengths)
        {
            int[] counts = new int[DeflateUtils.MaxSpecialLengthCode + 1];
            int extraBits = 0;
            int nonZeroSymbolLengths = DeflateUtils.MaxCode + 1;
            while (nonZeroSymbolLengths > 257 && symbolLengths[nonZeroSymbolLengths - 1] == 0) 
                --nonZeroSymbolLengths;

            int[] symbolLengthsCompact = DeflateUtils.CompactCodeLengths(symbolLengths, 0, nonZeroSymbolLengths);
            for (int i = 0; i < symbolLengthsCompact.Length; i++)
            {
                counts[symbolLengthsCompact[i]]++;
                if (symbolLengthsCompact[i] >= DeflateUtils.SpecialLengthCodeStart)
                {
                    extraBits += DeflateUtils.GetExtraBitsForLengthCode(symbolLengthsCompact[i]);
                    i++;
                }
            }

            int nonZeroDistanceLenghts = DeflateUtils.MaxDistanceCode + 1;
            while (nonZeroDistanceLenghts > 1 && symbolDistanceLengths[nonZeroDistanceLenghts - 1] == 0)
                --nonZeroDistanceLenghts;

            int[] symbolDistanceLengthsCompact = DeflateUtils.CompactCodeLengths(symbolDistanceLengths, 0, nonZeroDistanceLenghts);
            for (int i = 0; i < symbolDistanceLengthsCompact.Length; i++)
            {
                counts[symbolDistanceLengthsCompact[i]]++;
                if (symbolDistanceLengthsCompact[i] >= DeflateUtils.SpecialLengthCodeStart)
                {
                    extraBits += DeflateUtils.GetExtraBitsForLengthCode(symbolDistanceLengthsCompact[i]);
                    i++;
                }
            }

            int[] lengthsCodeLengths = DeflateUtils.GetCodeLengths(counts);
            int sum = DeflateUtils.CalculateLength(counts, lengthsCodeLengths);

            int[] map = DeflateUtils.LengthsMap;
            int nonZeroLengthCodes = DeflateUtils.MaxSpecialLengthCode + 1;
            while (nonZeroLengthCodes > 4 && lengthsCodeLengths[map[nonZeroLengthCodes - 1]] == 0)
                --nonZeroLengthCodes;

            int hclen = nonZeroLengthCodes - 4;

            return 14 + (hclen + 4) * 3 + sum + extraBits;
        }
    }
}
